from tensorflow import keras


class DayMinInputDecisionModel:
    def __init__(self,
                 learning_rate,
                 activation,
                 kernel_initializer,
                 epochs,
                 name,
                 day_shape,
                 min_shape
                 ):
        self.learning_rate = learning_rate
        self.activation = activation
        self.kernel_initializer = kernel_initializer
        self.epochs = epochs
        self.model = None
        self.name = name
        self.day_shape = day_shape
        self.min_shape = min_shape

    def __create_day_input(self, day_input):
        day_lstm = keras.layers.LSTM(units=16, activation=self.activation, name="day-lstm",
                                     kernel_initializer=self.kernel_initializer,
                                     return_sequences=True,
                                     dropout=0.02)(day_input)
        day_output = keras.layers.LayerNormalization(name="day-ln")(day_lstm)
        day_output = keras.layers.Dense(units=8,
                                        name="day-dense-1",
                                        activation=self.activation,
                                        kernel_initializer=self.kernel_initializer)(day_output)
        return day_output

    def __create_min_input(self, min_input):
        min_lstm = keras.layers.LSTM(units=64, activation=self.activation, name="min-lstm",
                                     kernel_initializer=self.kernel_initializer,
                                     dropout=0.02,
                                     return_sequences=True)(min_input)
        min_output = keras.layers.LayerNormalization(name="min-ln")(min_lstm)
        # min_output = keras.layers.LSTM(units=64, activation=self.activation, name="min-lstm-1",
        #                                kernel_initializer=self.kernel_initializer,
        #                                return_sequences=True)(min_output)
        min_output = keras.layers.Dropout(0.02)(min_output)
        min_output = keras.layers.Dense(units=8,
                                        name="min-dense-1",
                                        activation=self.activation,
                                        kernel_initializer=self.kernel_initializer)(min_output)
        return min_output

    def create_model(self):
        day_input = keras.layers.Input(shape=self.day_shape, name="day-input")
        min_input = keras.layers.Input(shape=self.min_shape, name="min-input")
        day_model = self.__create_day_input(day_input)
        min_model = self.__create_min_input(min_input)

        output = keras.layers.concatenate(inputs=[day_model, min_model], name="concat",axis=1)
        output = keras.layers.LSTM(units=8, activation=self.activation,
                                   name="concat-lstm",
                                   dropout=0.02,
                                   kernel_initializer=self.kernel_initializer)(output)
        output = keras.layers.Dense(units=72,
                                    name="output_dense",
                                    activation=self.activation,
                                    kernel_initializer=self.kernel_initializer)(output)

        output = keras.layers.LayerNormalization()(output)
        output = keras.layers.Dense(units=3,
                                    name="output",
                                    activation=keras.activations.softmax)(output)
        model = keras.Model(outputs=output,
                            name=self.name,
                            inputs=[day_input, min_input])

        model.compile(optimizer=keras.optimizers.Adam(learning_rate=self.learning_rate, weight_decay=0.01),
                      loss=keras.losses.categorical_crossentropy, metrics=["accuracy"])

        self.model = model
        return model

    def train_model(self,
                    train_data,
                    train_target,
                    val_data,
                    val_target):
        path_checkpoint = f'./model/{self.model.name}.h5'
        model_checkpoint_callback = keras.callbacks.ModelCheckpoint(
            monitor="val_loss",
            filepath=path_checkpoint,
            verbose=1,
            save_weights_only=True,
            save_best_only=True,
        )
        history = self.model.fit(
            x=train_data,
            y=train_target,
            batch_size=1024,
            epochs=self.epochs,
            verbose=1,
            callbacks=model_checkpoint_callback,
            validation_data=(val_data, val_target),
        )
        return history
